import Data.List (sortOn, intersect, (\\))

import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Utils (tok, thi3)

type Coord3     = (Int, Int, Int)
type Block      = (Coord3, Coord3)

type HeightMap  = Map Coord2 Height
type Coord2     = (Int, Int)
type Height     = Int

type BlockSet   = Set Block
type SupportMap = Map Block [Block]

parseLine :: String -> Block
parseLine s
    = ((x1, y1, z1), (x2, y2, z2))
    where
        toks         = map read . tok ",~" $ s
        (x1, y1, z1) = (toks !! 0, toks !! 1, toks !! 2)
        (x2, y2, z2) = (toks !! 3, toks !! 4, toks !! 5)

sortBlocks :: [Block] -> [Block]
sortBlocks = sortOn (thi3 . fst)

shadow :: Block -> [Coord2]
shadow ((x1, y1, _), (x2, y2, _)) = 
    [(x, y) | x <- [x1..x2], y <- [y1..y2]]

fall :: HeightMap -> Block -> (HeightMap, Block)
fall hm b@((x1, y1, z1), (x2, y2, z2)) = 
    (hm', b')
    where
        b'  = ((x1, y1, z1'), (x2, y2, z2'))
        z1' = (+1) . maximum . map (hm Map.!) . shadow $ b
        z2' = z1' + z2 - z1
        hm' = Map.union (Map.fromList (zip (shadow b) (repeat z2'))) hm

fallAll :: HeightMap -> [Block] -> [Block]
fallAll hm bs = 
    go hm bs []
    where
        go :: HeightMap -> [Block] -> [Block] -> [Block]
        go hm []     fallens = fallens
        go hm (b:bs) fallens = 
            go hm' bs (fallens ++ [b'])
            where
                (hm', b') = fall hm b

aboves :: [Block] -> Block -> [Block]
aboves bs b@(_, (_, _, z2)) = 
    bs'
    where
        bs' = filter (\x -> (shadow x `intersect` shadow b) /= [])
            . filter (\((_, _, zmin), _) -> zmin == z2 + 1)
            $ bs

belows :: [Block] -> Block -> [Block]
belows bs b@((_, _, z1), _) = 
    bs'
    where
        bs' = filter (\x -> (shadow x `intersect` shadow b) /= [])
            . filter (\(_, (_, _, zmax)) -> zmax == z1 - 1)
            $ bs

removable :: [Block] -> Block -> Bool
removable bs b = 
    all (\x -> ((belows bs x) \\ [b]) /= []) blocksAbove
    where
        blocksAbove = aboves bs b 

supports :: [Block] -> Block -> [Block]
supports bs ((_, _, 1), (_, _, _)) 
    = [((-500, -500, 0), (500, 500, 0))]
supports bs b = belows bs b

removeSome :: BlockSet -> SupportMap -> [Block] -> BlockSet
removeSome blockset _          []     = blockset
removeSome blockset supportmap blocks = 
    removeSome blockset' supportmap'' blocks'
    where
        blockset'    = blockset `Set.difference` (Set.fromList blocks)
        supportmap'  = Map.map (filter (`notElem` blocks)) supportmap
        supportmap'' = Map.filter (not . null) supportmap'
        blocks'      = map fst . Map.toList . Map.filter null $ supportmap'

countDownstream :: BlockSet -> SupportMap -> Block -> Int
countDownstream bs sm b = l - l' - 1
  where
    l  = length bs
    l' = length $ removeSome bs sm [b]

main = do
    input <- sortBlocks . map parseLine . lines <$> readFile "input.txt"

    let heightMap = Map.fromList [((x, y), 0) | x <- [0..500], y <- [0..500]]
    
    let blocks = fallAll heightMap $ input

    print $ length . filter (removable blocks) $ blocks

    -- For performance, we store *all* of the blocks in a set.
    let blockSet = Set.fromList blocks

    -- We only need to examine those bricks that are "unsafe", i.e.
    -- every brick that will trigger other bricks to fall if we remove it.
    let unsafes = filter (not . removable blocks) $ blocks

    -- Next, for every brick we store by which bricks it is supported.
    let supportMap = Map.fromList $ zip blocks (map (supports blocks) blocks)

    -- Finally, we just have to sum the result of every unsafe brick.
    print $ sum $ map (countDownstream blockSet supportMap) unsafes
